import asyncio
from mcp import ClientSession, StdioServerParameters, types
from mcp.client.stdio import stdio_client
from mcp.server.fastmcp import FastMCP

# 创建一个FastMCP服务器实例
mcp = FastMCP("MCP Server")


@mcp.tool()
def hello() -> str:
    return "hello mcp"


# 注册一个资源，URI为mcp//info，返回服务器的基本信息
@mcp.resource("mcp://info")
async def get_info():
    return {"name": "hello mcp", "version": "1.0.0", "description": "这是一个MCP服务器"}


@mcp.prompt("welcome")
async def get_welcome():
    return "welcome"


async def run_client():
    server_params = StdioServerParameters(command="python", args=[__file__, "serve"])
    # 开启一个新的子进程通过server_params，相当于开启了一个服务器的子进程,得到可读流和可写法
    async with stdio_client(server_params) as (read, write):
        # 基于读写流创建客户端会话
        async with ClientSession(read, write) as session:
            # 初始化连接
            await session.initialize()
            tools_result = await session.list_tools()
            print("[Tools]", [t.name for t in tools_result.tools])

            hello_result = await session.call_tool("hello", arguments={})
            tool_texts = []
            for block in hello_result.content:
                if isinstance(block, types.TextContent):
                    tool_texts.append(block.text)
            print(
                f"[Call hello] {"|".join(tool_texts) or str(hello_result.structuredContent)}"
            )

            resources_result = await session.list_resources()
            print(
                "[Resources]", [resource.uri for resource in resources_result.resources]
            )
            resource_texts = []
            for resource in resources_result.resources:
                resource_result = await session.read_resource(resource.uri)
                for block in resource_result.contents:
                    if isinstance(block, types.TextResourceContents):
                        resource_texts.append(block.text)
                print(f"[Resource:{resource.uri}] {"|".join(resource_texts)}")

            prompts_result = await session.list_prompts()
            print("[Prompts]", [prompt.name for prompt in prompts_result.prompts])
            prompt_texts = []
            for prompt in prompts_result.prompts:
                prompt_result = await session.get_prompt(prompt.name)
                for message in prompt_result.messages:
                    if isinstance(message, types.PromptMessage):
                        prompt_texts.append(message.content.text)
                print(f"[Prompt:{prompt.name}] {"|".join(prompt_texts)}")


def run_server_stdio():
    mcp.run(transport="stdio")


def main():
    import sys

    if len(sys.argv) >= 2 and sys.argv[1] == "serve":
        # 启动服务器模式
        run_server_stdio()
    else:
        asyncio.run(run_client())


if __name__ == "__main__":
    main()
